function [nnidx, dists] = slfindnn(Xr, Xq, k, varargin)
%SLFINDNN Finds the nearest neighbors
%
% [ Syntax ]
%   - nnidx = slfindnn(Xr, q, k, ...)
%   - nnidx = slfindnn(Xr, Xq, k, ...)
%   - nnidx = slfindnn(X, [], k, ...)
%   - [nnidx, dists] = slfindnn(Xr, Xq, k, ...)
%
% [ Arguments ]
%   - Xr:       the set of reference points (d x n matrix)
%   - Xq:       the set of query points (d x nq matrix)
%   - q:        a query point (d x 1 column vector)
%   - k:        the number of neighbors to search
%   
%   - nnidx:    the indices of neighbors
%   - dists:    the distances between the neighbors and the query point
%
% [ Description ]
%   - nnidx = slfindnn(Xr, q, k, ...) finds k nearest neighbors of q in the
%     set of reference points in Xr. 
%
%     Suppose we are to deal with d-dimensional points, and have n
%     reference points, then Xr should be a d x n real matrix, with each
%     column representing a point. The query point should be given in form
%     of a d x 1 column vector.
%
%     In the output, nnidx is a k x 1 column vector of the indices of the 
%     neighbors. The neighbors are sorted in ascending order of distances
%     to the query point. 
%
%   - nnidx = slfindnn(Xr, Xq, k, ...) finds k nearest neighbors for a 
%     set of query points in Xq from the reference point set Xr.
%
%     The usage is similar to that with single query. Suppose there are
%     nq query points. Then Xq should be a matrix of d x nq, with each
%     column being a query point. Correspondingly, nnidx will be a k x nq
%     matrix. Each column in nnidx stores the indices of the neighbors of
%     the corresponding query points as follows
%       $ nnidx(:, i) = slfindnn(Xr, Xq(:, i), k, ...) $
%
%     However, performing query in a batch is much more efficient than
%     doing it one by one in implementation.
%
%   - nnidx = slfindnn(X, [], k, ...) finds k nearest neighbors for each
%     point in X from the same set of points X. For each query point, the
%     point itself is excluded from its neighbor set.
%
%     If you intend to find nearest neighbors within the same set while
%     keeping the query point in its own neighbor set, you can simply write
%     as
%        $ nnidx = slfindnn(X, X, k, ... $
%  
%   - [nnidx, dists] = slfindnn(Xr, Xq, k, ...) additionally returns the 
%     corresponding distance values. 
%
%     Suppose Xq has nq points. Then dists is a k x nq matrix. As
%     introduced above, nnidx(p, i) is the index of the p-th neighbor of
%     the i-th query point. Then dists(p, i) is the corresponding distance
%     value between the the i-th query point and its p-th neighbor.
%
%     By default, the distances are Euclidean distance. You can also
%     specify other metric by setting the metric option (discussed in 
%     the following) as the distance used in searching.
%
%     You can further specify the following options
%     \{:
%         - metric:    The metric used in measuring the distance between 
%                      points. (default = 'eucdist')
%                           
%                      The metric can be either a string indicating the 
%                      name of the pre-defined metric (refer to the function 
%                      slmetric_pw for available predefined metrics).
%
%                      The metric can also be a function handle that supports 
%                      the following syntax
%                            $ D = f(X1, X2) $
%                      If X1 and X2 are d x n1 and d x n2 matrices respectively, 
%                      then D is an n1 x n2 matrix of pairwise metric values.
%
%         - radius:    The neighbor search radius, i.e. the maximum distances 
%                      between the query point and its neighbors. 
%
%                      By default, the option radius is set to 0, which means 
%                      no constraint on the distances.
%
%                      If radius is set to a positive value, then only and 
%                      points whose distances to the query point are less than 
%                      the radius can be considered as neighbors. 
%
%                      With a distance constraint, it is possible that for some 
%                      query points, the number of neighbors satisfying the 
%                      constaint is less than k. In these cases, the last entries 
%                      in nnidx are set to zeros. 
%
%                      For example, for thq query Xq(:,i), only r < k neighbors 
%                      are within the specified range, then nnidx(1:r, i) stores 
%                      the indices of those r neighbors, while nnidx(r+1:end, i) 
%                      are all zeros. If dists are also returned, the corresponding
%                      distance values are set to inf.
%
%         - dcache:    The size of distance cache in computation, i.e. the
%                      maximum number of distance values that can be stored
%                      in cache. (default = 1e7)
%
%                      In computing the distances between two large sets of
%                      samples, the number of pairwise distance values is
%                      very large, and thus it is probable that they can
%                      not be all accommodated in the memory. Typically, when
%                      memory is sufficient, using larger cache facilitates
%                      vectorization and thus is more efficient. 
%     \:}
%
% [ Remarks ]
%   - It is recommended to gather multiple queries together and invoke this
%     function to a batch of queries, which is much more efficient than
%     performing the query one by one.
%
%   - This function computes all pairwise-distances and then finds the
%     neighbors. For low-dimensional points, it is more efficient to use
%     kd-tree or its variants to locate the neighbors. 
%
%     The ANN MATLAB Wrapper is recommended for kd-tree or bd-tree based
%     nearest neighbor searching. It is available in MATLAB Central File
%     Exchange.
%
% [ History ]
%   - Created by Dahua Lin, on Jul 7, 2007
%

%% parse and verify input arguments

% Xq and Xr
error(nargchk(3, inf, nargin));
assert(isnumeric(Xr) && ndims(Xr) == 2, ...
    'sltoolbox:slfindnn:invalidarg', 'Xr should be a numeric matrix');

if ~isempty(Xq)
    assert(isnumeric(Xq) && ndims(Xq) == 2, ...
        'sltoolbox:slfindnn:invalidarg', 'Xq should be a numeric matrix');       
    assert(size(Xr, 1) == size(Xq, 1), ...
        'sltoolbox:slfindnn:dimmismatch', 'The point dimensions in Xq and Xr are different');
    exclude_self = false;
else
    Xq = Xr;
    exclude_self = true;
end
nr = size(Xr, 2);
nq = size(Xq, 2);

% k
assert(isnumeric(k) && isscalar(k) && k > 0 && k == fix(k), ...
    'sltoolbox:slfindnn:invalidarg', 'k should be an positive integer scalar');
if ~exclude_self
    assert(k <= nr, 'sltoolbox:slfindnn:k_toolarge', ...
        'k should not be larger than the number of reference points');
else
    assert(k < nr, 'sltoolbox:slfindnn:k_toolarge', ...
        'k should be less than the number of points in self-exclusion mode');
end

% options
opts = struct( ...
    'metric', 'eucdist', ...
    'radius', 0, ...
    'dcache', 1e7);
opts = setopts(opts, varargin{:});

assert(ischar(opts.metric) || isa(opts.metric, 'function_handle'), ...
    'sltoolbox:slfindnn:invalidopt', 'The option metric should be a string or a function handle.');
assert(isfloat(opts.radius) && isscalar(opts.radius) && opts.radius >= 0, ...
    'sltoolbox:slfindnn:invalidopt', 'The option radius should be a non-negative scalar value.');
assert(isnumeric(opts.dcache) && isscalar(opts.dcache) && opts.dcache > 0, ...
    'sltoolbox:slfindnn:invalidopt', 'The option dcache should be a positive scalar value.');
assert(opts.dcache >= nr, 'sltoolbox:slfindnn:cache_toosmall', ...
    'The option dcache is too small to hold distances for one query');


%% main

if ischar(opts.metric)
    fm = @(x, y) slmetric_pw(x, y, opts.metric);
else
    fm = opts.metric;
end
r = opts.radius;

if nr * nq <= opts.dcache  % single batch
    if nargout < 2
        nnidx = core_findnn(Xr, Xq, fm, k, r, exclude_self);
    else
        [nnidx, dists] = core_findnn(Xr, Xq, fm, k, r, exclude_self);
    end
    
else  % multiple batches
    
    % divide batch
    bn = floor(opts.dcache / nr);
    sp = 1:bn:nq;
    ep = [sp(2:end) - 1, nq];
    nb = length(sp);
    
    if nargout < 2
        nnidx = zeros(k, nq);    
        for i = 1 : nb
            nnidx(:, sp(i):ep(i)) = ...
                core_findnn(Xr, Xq(:, sp(i):ep(i)), fm, k, r, exclude_self);
        end
    else
        nnidx = zeros(k, nq);
        dists = zeros(k, nq);
        for i = 1 : nb
            [nnidx(:, sp(i):ep(i)), dists(:, sp(i):ep(i))] = ...
                core_findnn(Xr, Xq(:, sp(i):ep(i)), fm, k, r, exclude_self);
        end
    end
end


%% core function

function [nnidx, dists] = core_findnn(Xr, Xq, fm, k, r, exc_self)

dists = fm(Xr, Xq);
if exc_self
    n = size(Xr, 2);
    dists(1:(n+1):(n*n)) = inf;
end
[dists, nnidx] = sort(dists, 1);


dists = dists(1:k, :);
nnidx = nnidx(1:k, :);

if r > 0
    is_exceed = (dists > r);
    dists(is_exceed) = inf;
    nnidx(is_exceed) = 0;
end

